import os
import sys
import numpy as np
from multiprocessing import Pool
from frb_pipeline_tools import (
    extract_candidate_segment,
    normalise,
    dedisperse,
    plot_frb_four_panel_summary
)

def process_one_candidate(args):
    """
    处理一条候选体，提取数据、处理并绘图。
    """
    line, line_number, fits_list, fits_index_map, npz_dir, plot_output_dir, save_npz = args

    try:
        parts = line.strip().split()
        fits_name = parts[0]
        sample_id = int(parts[1])
        DM = float(parts[2])
        snr = float(parts[-1])

        if fits_name not in fits_index_map:
            print(f"[Skipped] FITS not found: {fits_name}")
            return

        index = fits_index_map[fits_name]
        fits_path = fits_list[index]

        data, freqs, times, tsamp, mjd_start = extract_candidate_segment(
            fits_list, index, sample_id, DM
        )
        data_clean = normalise(data)
        data_dedisp = dedisperse(data_clean, freqs, tsamp, DM)

        base = os.path.splitext(fits_name)[0]
        out_png = os.path.join(
            plot_output_dir,
            f"{base}_L{line_number}_DM{int(DM)}.png"
        )

        plot_frb_four_panel_summary(
            data_raw=data_clean,
            data_dedispersed=data_dedisp,
            freqs=freqs,
            times=times,
            tsamp=tsamp,
            snr=snr,
            dm=DM,
            sample_id=sample_id,
            fits_file=fits_path,
            mjd_start=mjd_start,
            line_number=line_number,
            outname=out_png
        )

        if save_npz:
            npz_path = os.path.join(npz_dir, base + ".npz")
            np.savez(npz_path, data=data, freqs=freqs, times=times,
                     data_dedisp=data_dedisp)

    except Exception as e:
        print(f"[Error] Failed sample {sample_id} in {fits_name}: {e}")


if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: python run_frb_search.py <fits_dir> <filtered.txt> <plot_output_dir> [--save_npz]")
        sys.exit(1)

    fits_dir = sys.argv[1]
    cand_txt_path = sys.argv[2]
    plot_output_dir = sys.argv[3]
    save_npz = "--save_npz" in sys.argv

    npz_dir = os.path.join(os.path.dirname(plot_output_dir), "npz")
    os.makedirs(plot_output_dir, exist_ok=True)
    if save_npz:
        os.makedirs(npz_dir, exist_ok=True)

    fits_list = sorted([
        os.path.join(fits_dir, f) for f in os.listdir(fits_dir)
        if f.endswith(".fits")
    ])
    fits_index_map = {os.path.basename(f): i for i, f in enumerate(fits_list)}

    with open(cand_txt_path, 'r') as f:
        lines = [line.strip() for line in f if line.strip()]

    core = 15
    print(f"Total candidates: {len(lines)}")
    print(f"Using {core} CPU cores...")

    args_list = [
        (line, idx + 1, fits_list, fits_index_map, npz_dir, plot_output_dir, save_npz)
        for idx, line in enumerate(lines)
    ]
    try:
        with Pool(processes=core) as pool:
            pool.map(process_one_candidate, args_list)
    except KeyboardInterrupt:
        print("🔴 Interrupted by user, terminating...")
        pool.terminate()
        pool.join()

    print("✅ All candidates processed.")